from transformers import AutoTokenizer, AutoModelForTokenClassification
from transformers import pipeline
import torch
from pprint import pprint
import matplotlib.pyplot as plt
import seaborn as sns

# Check if GPU is available
device = 0 if torch.cuda.is_available() else -1
print(f"Using device: {'GPU' if device == 0 else 'CPU'}")

# Initialize tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("dslim/bert-large-NER")
model = AutoModelForTokenClassification.from_pretrained("dslim/bert-large-NER")

# Create NER pipeline
ner_pipeline = pipeline(
    "ner",
    model=model,
    tokenizer=tokenizer,
    aggregation_strategy="simple",
    device=device,
)

# Fun examples showcasing different scenarios
examples = [
    # Famous fictional characters and locations
    "Harry Potter lives at 4 Privet Drive with the Dursleys.",
    # Historical figures and events
    "Albert Einstein published his theory of relativity in 1915 while working in Berlin.",
    # Modern tech and business
    "Mark Zuckerberg founded Facebook at Harvard before moving to Silicon Valley.",
    # Mixed entities (people, organizations, locations, dates)
    "The NASA mission to Mars began in July 2020, with Perseverance landing on February 18, 2021.",
    # Sports and entertainment
    "Michael Jordan led the Chicago Bulls to six NBA championships during the 1990s.",
    # Multiple entity types in one sentence
    "Apple CEO Tim Cook announced the new iPhone launch at their headquarters in Cupertino, California.",
]


def analyze_entities(text):
    """Analyze entities in text and return detailed results."""
    entities = ner_pipeline(text)
    return {
        "text": text,
        "entities": entities,
        "entity_count": len(entities),
        "entity_types": set(entity["entity_group"] for entity in entities),
    }


def visualize_confidence_scores(results):
    """Create a bar plot of entity confidence scores."""
    entities = [e["word"] for e in results["entities"]]
    scores = [e["score"] for e in results["entities"]]

    plt.figure(figsize=(10, 5))
    sns.barplot(x=scores, y=entities)
    plt.title("Entity Detection Confidence Scores")
    plt.xlabel("Confidence Score")
    plt.ylabel("Entity")
    plt.tight_layout()
    return plt


def display_analysis(analysis):
    """Display the analysis results in a readable format."""
    print("\n" + "=" * 80)
    print(f"Analyzing: '{analysis['text']}'")
    print("=" * 80)

    # Display entity count and types
    print(
        f"\nFound {analysis['entity_count']} entities of {len(analysis['entity_types'])} different types:"
    )
    print(f"Types: {', '.join(analysis['entity_types'])}")

    # Display detailed entity information
    print("\nDetailed Entity Analysis:")
    for entity in analysis["entities"]:
        print(f"\n• Entity: {entity['word']}")
        print(f"  - Type: {entity['entity_group']}")
        print(f"  - Confidence: {entity['score']:.4f}")

    # Create and save visualization
    plt = visualize_confidence_scores(analysis)
    output_path = f"entity_confidence_{examples.index(analysis['text'])+1}.png"
    plt.savefig(output_path)
    plt.close()
    print(f"\nConfidence score visualization saved as: {output_path}")


def main():
    print("\n🔍 Welcome to the HuggingFace NER Showcase! 🔍")
    print(
        "This demo will analyze various texts to demonstrate the capabilities of dslim/bert-large-NER"
    )

    # Process each example
    for i, example in enumerate(examples, 1):
        print(f"\n\nProcessing Example {i}/{len(examples)}...")
        results = analyze_entities(example)
        display_analysis(results)

    print("\n\n✨ Analysis Complete! ✨")
    print(
        "Check out the generated PNG files to see the confidence score visualizations!"
    )


if __name__ == "__main__":
    main()
